import os
import shutil

import numpy as np

from option import parse_args


def first_five_frames_exist(real_video_crop, fake_video_crop):
    for i in range(5):
        image_name = '{}_0.png'.format(i * 10)
        file_exist = os.path.exists(os.path.join(real_video_crop, image_name)) and \
                     os.path.exists(os.path.join(fake_video_crop, image_name))
        if not file_exist:
            return False
    return True


def copy_five_frame(sampled_videos, args):
    real_set = set()
    with open(os.path.join(args.dst_dir, 'real_fake_pairs.dat'), 'w') as f:
        for sampled_video in sampled_videos:
            sampled_video_split = sampled_video.split(',')
            assert len(sampled_video_split) == 3
            part = sampled_video_split[0]
            real = sampled_video_split[1]
            fake = sampled_video_split[2]
            real_video_crop = os.path.join(args.src_dir, 'crops', part, real)
            fake_video_crop = os.path.join(args.src_dir, 'crops', part, fake)
            if real in real_set:
                continue
            real_set.add(real)
            if not first_five_frames_exist(real_video_crop, fake_video_crop):
                continue

            for i in range(5):
                image_name = '{}_0.png'.format(i * 10)
                if not os.path.exists(os.path.join(real_video_crop, image_name)):
                    continue

                dst_real_image = f'{part}_{real}__{image_name}'
                if os.path.exists(os.path.join(args.src_dir, dst_real_image)):
                    continue
                dst_fake_image = f'{part}_{real}_{fake}__{image_name}'

                shutil.copy(os.path.join(real_video_crop, image_name), os.path.join(args.dst_dir, dst_real_image))
                shutil.copy(os.path.join(fake_video_crop, image_name), os.path.join(args.dst_dir, dst_fake_image))
                f.write('{},{}\n'.format(dst_real_image, dst_fake_image))


def get_video_names_by_mode(args, mode):
    video_name_set = []
    video_names = open(os.path.join(args.src_dir, '{}.txt'.format(mode)), 'r').readlines()
    for video_name in video_names:
        video_name_split = video_name.split(',')
        part = video_name_split[0]
        real = video_name_split[1]
        fake = video_name_split[2].strip()
        video_name_set.append('{},{},{}'.format(part, real, fake))
    return video_name_set


def main():
    args = parse_args()
    args.src_dir = os.path.join(args.src_dir, 'deepfake-detection-challenge')
    args.dst_dir = os.path.join(args.dst_dir, 'dfdc')
    os.makedirs(args.dst_dir, exist_ok=True)

    # manipulated_path = os.path.join(args.src_dir, 'Celeb-synthesis')
    videos = get_video_names_by_mode(args, 'test')
    np.random.seed(111)
    sampled_videos = np.random.choice(videos, 200)

    copy_five_frame(sampled_videos, args)


if __name__ == '__main__':
    main()
